import cv2
import os
from pathlib import Path

##################主要是对标签以及数据集的处理

def set_class():  ###调整标签
    img_list = os.listdir('train/masks')
    for file in img_list:
        img = cv2.imread(f'train/masks/{file}')
        img[img == 255] = 255
        img[img == 2] = 0
        img[img == 3] = 255
        cv2.imwrite(f'train/masks_color/{file}', img)


def sampling():  ###增加样本数据量
    img_dir = 'train/imgs'
    mask_dir = 'train/masks_color'
    save_dir = 'train/sampling_train'
    Path(save_dir).mkdir(parents=True, exist_ok=True)
    Path(os.path.join(save_dir, "imgs")).mkdir(parents=True, exist_ok=True)
    Path(os.path.join(save_dir, 'masks')).mkdir(parents=True, exist_ok=True)
    img_list = os.listdir(img_dir)
    for file in img_list:
        img = cv2.imread(os.path.join(img_dir, file))
        mask = cv2.imread(os.path.join(mask_dir, file[:-4] + ".png"))

        cv2.imwrite(os.path.join(save_dir, f'imgs/{file}'), img)
        cv2.imwrite(os.path.join(save_dir, f'masks/{file[:-4]}.png'), mask)

        img_ud = cv2.flip(img, flipCode=0)  # 上下翻转
        mask_ud = cv2.flip(mask, flipCode=0)
        cv2.imwrite(os.path.join(save_dir, f'imgs/ud_{file}'), img_ud)
        cv2.imwrite(os.path.join(save_dir, f'masks/ud_{file[:-4]}.png'), mask_ud)

        img_lr = cv2.flip(img, flipCode=1)  # 左右翻转
        mask_lr = cv2.flip(mask, flipCode=1)
        cv2.imwrite(os.path.join(save_dir, f'imgs/lr_{file}'), img_lr)
        cv2.imwrite(os.path.join(save_dir, f'masks/lr_{file[:-4]}.png'), mask_lr)

        img_udlr = cv2.flip(img, flipCode=-1)  # 上下+左右翻转
        mask_udlr = cv2.flip(mask, flipCode=-1)
        cv2.imwrite(os.path.join(save_dir, f'imgs/udlr_{file}'), img_udlr)
        cv2.imwrite(os.path.join(save_dir, f'masks/udlr_{file[:-4]}.png'), mask_udlr)

if __name__=='__main__':
    sampling()
